#include "denoiser.h"
#include "audio_utils.h"
#include <iostream>
#include <torch/script.h>
#include <torch/torch.h>

class DenoiserImpl : public torch::nn::Module {
    // 实现与Python模型相同的网络结构
};

Denoiser::Denoiser(const std::string& model_path, const std::string& device) 
    : device_(device == "cuda" && torch::cuda::is_available() ? 
             torch::kCUDA : torch::kCPU)
{
    try {
        // 加载模型
        model_ = torch::jit::load(model_path);
        
        // 将模型移到指定设备并设置为评估模式
        model_.to(device_);
        model_.eval();
    }
    catch (const c10::Error& e) {
        std::cerr << "Error loading model: " << e.what() << std::endl;
        throw;
    }
}

torch::Tensor Denoiser::denoise(const torch::Tensor& audio) {
    torch::NoGradGuard no_grad;
    
    const int segment_size = 16000;  // 1秒的音频
    const int overlap = segment_size / 2;  // 50% 重叠
    
    // 确保输入在正确的设备上
    auto input = audio.to(device_);
    
    // 计算总长度和段数
    int total_length = input.size(0);
    std::vector<torch::Tensor> processed_segments;
    
    // 处理每个段
    for (int start = 0; start < total_length; start += (segment_size - overlap)) {
        // 提取段
        int end = std::min(start + segment_size, total_length);
        auto segment = input.slice(0, start, end);
        
        // 如果是最后一段，需要填充
        if (segment.size(0) < segment_size) {
            auto padding_size = segment_size - segment.size(0);
            segment = torch::pad(segment, {0, padding_size});
        }
        
        // 重塑为 [1, 1, time] 并处理
        auto shaped_segment = segment.unsqueeze(0).unsqueeze(0);
        auto processed = model_.forward({shaped_segment}).toTensor();
        
        // 移除多余的维度
        processed = processed.squeeze(0).squeeze(0);
        
        // 如果不是第一段，移除重叠部分的前半部分
        if (!processed_segments.empty()) {
            processed = processed.slice(0, overlap);
        }
        
        // 如果是最后一段，移除填充
        if (end == total_length && segment.size(0) < segment_size) {
            processed = processed.slice(0, 0, segment.size(0));
        }
        
        processed_segments.push_back(processed);
    }
    
    // 连接所有处理后的段
    auto result = torch::cat(processed_segments);
    
    // 确保输出长度与输入相同
    if (result.size(0) > total_length) {
        result = result.slice(0, 0, total_length);
    }
    
    return result.to(torch::kCPU);
}

void Denoiser::process_file(const std::string& input_path, const std::string& output_path) {
    try {
        // 加载音频
        auto [audio_data, sr] = AudioUtils::load_audio(input_path);
        sample_rate_ = sr;  // 保存采样率
        
        std::cout << "Processing audio with " << audio_data.size() << " samples\n";
        
        // 转换为tensor
        torch::Tensor audio_tensor = torch::from_blob(
            audio_data.data(),
            {static_cast<long>(audio_data.size())},
            torch::kFloat32
        ).clone();  // 创建数据的副本
        
        // 降噪处理
        auto output_tensor = denoise(audio_tensor);
        
        // 确保输出长度正确
        if (output_tensor.numel() != audio_data.size()) {
            std::cout << "Warning: Output size (" << output_tensor.numel() 
                      << ") differs from input size (" << audio_data.size() << ")\n";
        }
        
        // 转换回vector
        std::vector<float> output_data(
            output_tensor.data_ptr<float>(),
            output_tensor.data_ptr<float>() + output_tensor.numel()
        );
        
        // 保存结果
        AudioUtils::save_audio(output_path, output_data, sample_rate_);
        
        std::cout << "Successfully processed audio:\n"
                  << "  Input samples: " << audio_data.size() << "\n"
                  << "  Output samples: " << output_data.size() << "\n"
                  << "  Sample rate: " << sample_rate_ << "\n";
    }
    catch (const std::exception& e) {
        std::cerr << "Error processing file: " << e.what() << std::endl;
        throw;
    }
} 